import os

from pycocotools.coco import COCO


def convert_coco_to_yolo(coco_json_path, output_folder):
    # 加载COCO标注文件
    coco = COCO(coco_json_path)
    # 获取所有图像的ID
    image_ids = coco.getImgIds()
    # 创建输出文件夹如果不存在
    os.makedirs(output_folder, exist_ok=True)

    for image_id in image_ids:
        image_info = coco.loadImgs(image_id)[0]
        image_width = image_info['width']
        image_height = image_info['height']

        # 获取属于该图像的所有标注信息
        ann_ids = coco.getAnnIds(imgIds=image_id, iscrowd=None)
        annotations = coco.loadAnns(ann_ids)

        # 创建YOLO格式的标注文件路径
        output_file = os.path.join(output_folder, f"{os.path.splitext(image_info['file_name'])[0]}.txt")
        with open(output_file, 'w') as f:
            for ann in annotations:
                bbox = ann['bbox']
                x_center = (bbox[0] + bbox[2] / 2) / image_width
                y_center = (bbox[1] + bbox[3] / 2) / image_height
                width = bbox[2] / image_width
                height = bbox[3] / image_height
                category_id = ann['category_id']
                # category = coco.loadCats(category_id)[0]['name']  # 获取类别名，可能需要根据你的实际需求调整这部分逻辑（例如，使用id代替名称）
                line = f"{category_id} {x_center} {y_center} {width} {height}\n"
                f.write(line)
    print("Conversion completed.")


if __name__ == '__main__':
    t = [('test.json', 'test'), ('train.json', 'train'), ('valid.json', 'valid')]
    root_dir = r'D:\project\py\datasets\worker_safety'
    for i in t:
        convert_coco_to_yolo(
            f'{root_dir}/annotations/{i[0]}',
            f'{root_dir}/labels/{i[1]}'
        )
